package com.larry.spark.sql

import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.sql.{Row, SparkSession}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, IntegerType, LongType, StructField, StructType}

object Sql_Oper_UDF_1 {

  def main(args: Array[String]): Unit = {
    //TODO  使用spark coalesce  缩减分区
    //TODO  默认情况下 缩减分区不会shuffle

    val conf = new SparkConf().setMaster("local[*]").setAppName("sql")

    //创建session对象
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()

    import spark.implicits._

    //读取json
    val df = spark.read.json("input/user.json")

    df.createOrReplaceTempView("user")

    val myAvgUDAF = new MyAvgUDAF

    spark.udf.register("myavg",myAvgUDAF)

    spark.sql("select myavg(age) from user").show()

    //关闭资源
    spark.stop()

  }
}
class MyAvgUDAF extends UserDefinedAggregateFunction{

  //输入
  override def inputSchema: StructType = StructType(Array(StructField("age",IntegerType)))

  override def bufferSchema: StructType = {
    StructType(Array(StructField("sum",LongType),StructField("count",LongType)))
  }

  //输出数据类型
  override def dataType: DataType = DoubleType

  override def deterministic: Boolean = true

  //缓冲区初始化
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0L
    buffer(1) = 0L
  }

  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    if (!input.isNullAt(0)){
      buffer(0) = buffer.getLong(0) + input.getInt(0)
      buffer(1) = buffer.getLong(1) + 1
    }
  }

  //合并
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  //输出计算
  override def evaluate(buffer: Row): Any = {
    buffer.getLong(0).toDouble / buffer.getLong(1)
  }
}